"""
Hi-C training utilities adapted from the original train_util.py
Specialized for Hi-C temporal diffusion model training
"""
import copy
import functools
import os

import torch as th
import torch.distributed as dist
from torch.nn.parallel.distributed import DistributedDataParallel as DDP
from torch.optim import AdamW

from . import dist_util, logger
from .fp16_util import MixedPrecisionTrainer
from .nn import update_ema
from .resample import LossAwareSampler, UniformSampler
import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import pearsonr, spearmanr

INITIAL_LOG_LOSS_SCALE = 20.0

def plot_hic_loss(losses, args_train):
    """Plot training loss for Hi-C model"""
    losses_np = np.array([i.detach().cpu().numpy() if th.is_tensor(i) else i for i in losses])
    
    # Calculate moving average
    window_size = max(int(len(losses_np) / 100), 10)
    if len(losses_np) >= window_size:
        windowed_mean_loss = np.convolve(losses_np, np.ones(window_size) / window_size, mode='valid')
        windowed_x_vals = np.arange(window_size - 1, len(losses_np))
    else:
        windowed_mean_loss = losses_np
        windowed_x_vals = np.arange(len(losses_np))
    
    # Plotting
    fig, ax = plt.subplots(figsize=(10, 6), dpi=150)
    x_vals = np.arange(len(losses_np))
    
    plt.plot(x_vals, losses_np, label='Training Loss', alpha=0.3, color='blue')
    plt.plot(windowed_x_vals, windowed_mean_loss, label=f'Moving Average (window={window_size})', 
             color='red', linewidth=2)
    
    plt.xlabel('Training Steps')
    plt.ylabel('Loss')
    plt.title('Hi-C Diffusion Model Training Loss')
    plt.legend()
    plt.grid(True, alpha=0.3)
    
    # Save plot
    save_dir = args_train.get('logger_path', '.')
    plt.savefig(os.path.join(save_dir, 'hic_training_loss.png'), 
                dpi=150, bbox_inches='tight', transparent=True)
    plt.savefig(os.path.join(save_dir, 'hic_training_loss.pdf'), 
                bbox_inches='tight', transparent=True)
    plt.close()

def calculate_hic_metrics(pred_matrix, true_matrix):
    """
    Calculate Hi-C specific evaluation metrics
    
    Args:
        pred_matrix: predicted Hi-C matrix [H, W] 
        true_matrix: ground truth Hi-C matrix [H, W]
    
    Returns:
        dict of metrics
    """
    # Flatten matrices for correlation calculation
    pred_flat = pred_matrix.flatten()
    true_flat = true_matrix.flatten()
    
    # Remove zero entries for correlation calculation
    mask = (true_flat > 0) | (pred_flat > 0)
    if mask.sum() == 0:
        return {'pearson': 0, 'spearman': 0, 'mse': float('inf')}
    
    pred_masked = pred_flat[mask]
    true_masked = true_flat[mask]
    
    # Calculate correlations
    try:
        pearson_corr, _ = pearsonr(pred_masked, true_masked)
        spearman_corr, _ = spearmanr(pred_masked, true_masked)
    except:
        pearson_corr = 0
        spearman_corr = 0
    
    # Calculate MSE
    mse = np.mean((pred_masked - true_masked) ** 2)
    
    return {
        'pearson': pearson_corr,
        'spearman': spearman_corr, 
        'mse': mse
    }

class HiCTrainLoop:
    """Training loop for Hi-C diffusion model"""
    
    def __init__(
        self,
        *,
        model,
        diffusion,
        data,
        batch_size,
        microbatch,
        lr,
        ema_rate,
        log_interval,
        save_interval,
        resume_checkpoint,
        use_fp16=False,
        fp16_scale_growth=1e-3,
        schedule_sampler=None,
        weight_decay=0.0,
        lr_anneal_steps=0,
        multi_t_training=False,
        num_t_per_batch=1,
        t_multi_strategy="linspace",
        t_stride=0,
    ):
        self.model = model
        # Ensure model parameters are on the correct device **before** any wrapper (FP16 or DDP)
        if th.cuda.is_available():
            self.model = self.model.to(th.device("cuda"))
        self.diffusion = diffusion
        self.data = data
        self.batch_size = batch_size
        self.microbatch = microbatch if microbatch > 0 else batch_size
        self.lr = lr
        self.ema_rate = (
            [ema_rate]
            if isinstance(ema_rate, float)
            else [float(x) for x in ema_rate.split(",")]
        )
        self.log_interval = log_interval
        self.save_interval = save_interval
        self.resume_checkpoint = resume_checkpoint
        self.use_fp16 = use_fp16
        self.fp16_scale_growth = fp16_scale_growth
        self.schedule_sampler = schedule_sampler or UniformSampler(diffusion)
        self.weight_decay = weight_decay
        self.lr_anneal_steps = lr_anneal_steps
        # Multi-timestep training options
        self.multi_t_training = bool(multi_t_training)
        self.num_t_per_batch = int(num_t_per_batch) if num_t_per_batch is not None else 1
        self.t_multi_strategy = t_multi_strategy
        self.t_stride = int(t_stride) if t_stride is not None else 0

        self.step = 0
        self.resume_step = 0
        self.global_batch = self.batch_size
        self.sync_cuda = th.cuda.is_available()
        self.loss_list = []
        
        # Setup mixed precision trainer
        self.mp_trainer = MixedPrecisionTrainer(
            model=self.model,
            use_fp16=self.use_fp16,
            fp16_scale_growth=fp16_scale_growth,
        )

        # Setup optimizer
        self.opt = AdamW(
            self.mp_trainer.master_params, lr=self.lr, weight_decay=self.weight_decay
        )
        
        # Load checkpoint if resuming
        if self.resume_step:
            self._load_optimizer_state()
            self.ema_params = [
                self._load_ema_parameters(rate) for rate in self.ema_rate
            ]
        else:
            self.ema_params = [
                copy.deepcopy(self.mp_trainer.master_params)
                for _ in range(len(self.ema_rate))
            ]

        # Setup distributed training
        if th.cuda.is_available() and dist.is_initialized():
            self.use_ddp = True
            self.ddp_model = DDP(
                self.model,
                device_ids=[dist_util.dev()],
                output_device=dist_util.dev(),
                broadcast_buffers=False,
                bucket_cap_mb=128,
                find_unused_parameters=True,
            )
        else:
            self.use_ddp = False
            self.ddp_model = self.model

    def _load_ema_parameters(self, rate):
        """Load EMA parameters from checkpoint"""
        ema_params = copy.deepcopy(self.mp_trainer.master_params)
        main_checkpoint = self.resume_checkpoint
        ema_checkpoint = self._find_ema_checkpoint(main_checkpoint, self.resume_step, rate)
        if ema_checkpoint:
            logger.log(f"Loading EMA from checkpoint: {ema_checkpoint}")
            state_dict = dist_util.load_state_dict(
                ema_checkpoint, map_location=dist_util.dev()
            )
            ema_params = self.mp_trainer.state_dict_to_master_params(state_dict)
        dist_util.sync_params(ema_params)
        return ema_params

    def _load_optimizer_state(self):
        """Load optimizer state from checkpoint"""
        main_checkpoint = self.resume_checkpoint
        opt_checkpoint = os.path.join(
            os.path.dirname(main_checkpoint), f"opt{self.resume_step:06}.pt"
        )
        if os.path.exists(opt_checkpoint):
            logger.log(f"Loading optimizer state from checkpoint: {opt_checkpoint}")
            state_dict = dist_util.load_state_dict(
                opt_checkpoint, map_location=dist_util.dev()
            )
            self.opt.load_state_dict(state_dict)

    def _find_ema_checkpoint(self, main_checkpoint, step, rate):
        """Find EMA checkpoint file"""
        if main_checkpoint is None:
            return None
        filename = f"ema_{rate}_{step:06d}.pt"
        path = os.path.join(os.path.dirname(main_checkpoint), filename)
        if os.path.exists(path):
            return path
        return None

    def run_loop(self):
        """Main training loop"""
        data_iter = iter(self.data)
        
        while (
            not self.lr_anneal_steps
            or self.step + self.resume_step < self.lr_anneal_steps
        ):
            try:
                batch = next(data_iter)
            except StopIteration:
                data_iter = iter(self.data)
                batch = next(data_iter)

            self.run_step(batch)
            
            if self.step % self.log_interval == 0:
                logger.dumpkvs()
            
            if self.step % self.save_interval == 0:
                self.save()
                # Stop if testing
                if os.environ.get("DIFFUSION_TRAINING_TEST", "") and self.step > 0:
                    return
                    
            self.step += 1
            
        # Save final checkpoint
        if (self.step - 1) % self.save_interval != 0:
            self.save()

    def run_step(self, batch):
        """Run a single training step"""
        self.forward_backward(batch)
        took_step = self.mp_trainer.optimize(self.opt)
        if took_step:
            self._update_ema()
        self._anneal_lr()
        self.log_step()

    def _select_t_list(self):
        """Return a list of integer timesteps for multi-t training according to strategy/stride."""
        K = min(self.num_t_per_batch, self.diffusion.num_timesteps)
        if self.t_multi_strategy == "random":
            if self.t_stride and self.t_stride > 0:
                grid = np.arange(0, self.diffusion.num_timesteps, self.t_stride, dtype=int)
                K = min(K, len(grid))
                t_list = np.random.choice(grid, size=K, replace=False)
            else:
                t_list = np.random.choice(self.diffusion.num_timesteps, size=K, replace=False)
            t_list.sort()
        else:
            if self.t_stride and self.t_stride > 0:
                t_list = np.arange(0, self.diffusion.num_timesteps, self.t_stride, dtype=int)[:K]
            else:
                t_list = np.linspace(0, self.diffusion.num_timesteps - 1, K, dtype=int)
        return [int(x) for x in t_list.tolist()] if hasattr(t_list, 'tolist') else [int(x) for x in t_list]

    def _train_microbatch(self, micro_hic, micro_bio_time, last_batch, extra_loss_fn=None):
        """
        Compute loss for one microbatch, optionally over multiple timesteps, and backprop.
        extra_loss_fn: callable(micro_hic) -> extra_loss_tensor or None; added into total loss when provided.
        Returns the scalar tensor used for backward (aggregated loss).
        """
        # Choose timesteps
        steps = []
        if not self.multi_t_training or self.num_t_per_batch <= 1:
            # use sampler
            t, weights = self.schedule_sampler.sample(micro_hic.shape[0], dist_util.dev())
            steps.append((t, weights))
        else:
            t_list = self._select_t_list()
            for t_scalar in t_list:
                t = th.full((micro_hic.shape[0],), int(t_scalar), device=dist_util.dev(), dtype=th.long)
                weights = th.ones_like(t, dtype=th.float32)
                steps.append((t, weights))

        agg_loss = 0.0
        for idx, (t, weights) in enumerate(steps):
            compute_losses = functools.partial(
                self.diffusion.training_losses,
                self.ddp_model,
                micro_hic,
                t,
                model_kwargs={'bio_time': micro_bio_time},
            )
            if last_batch or not self.use_ddp:
                losses = compute_losses()
            else:
                with self.ddp_model.no_sync():
                    losses = compute_losses()

            # Optional extra regularization
            if extra_loss_fn is not None:
                extra = extra_loss_fn(micro_hic)
                try:
                    if extra is not None and extra.item() > 0:
                        losses["extra_reg"] = extra
                        losses["loss"] = losses["loss"] + losses["extra_reg"]
                except Exception:
                    pass

            if isinstance(self.schedule_sampler, LossAwareSampler):
                self.schedule_sampler.update_with_local_losses(t, losses["loss"].detach())

            step_loss = (losses["loss"] * weights).mean()
            # Average when multiple timesteps
            denom = float(len(steps))
            agg_loss = agg_loss + step_loss / denom
            # 当一次microbatch包含多个t时，日志按平均尺度记录，避免放大
            scale = (1.0 / denom) if denom > 1 else 1.0
            log_loss_dict(self.diffusion, t, {k: (v * weights * scale) for k, v in losses.items()})

        self.mp_trainer.backward(agg_loss)
        if (self.step + self.resume_step) < 2:   # 只看前两步
            unused = [n for n, p in self.ddp_model.module.named_parameters()
                        if p.requires_grad and (p.grad is None)]
            if unused:
                print("本步未参与反传的参数:", unused)
        
        return agg_loss

    def forward_backward(self, batch):
        """Forward and backward pass"""
        self.mp_trainer.zero_grad()
        hic_data = batch['hic'].to(dist_util.dev())
        bio_time = batch['time'].to(dist_util.dev())
        last_tracked = None
        for i in range(0, hic_data.shape[0], self.microbatch):
            micro_hic = hic_data[i : i + self.microbatch]
            micro_bio_time = bio_time[i : i + self.microbatch]
            last_batch = (i + self.microbatch) >= hic_data.shape[0]
            last_tracked = self._train_microbatch(micro_hic, micro_bio_time, last_batch)
        if last_tracked is not None:
            self.loss_list.append(last_tracked.detach().cpu())

    def _update_ema(self):
        """Update EMA parameters"""
        for rate, params in zip(self.ema_rate, self.ema_params):
            update_ema(params, self.mp_trainer.master_params, rate=rate)

    def _anneal_lr(self):
        """Anneal learning rate"""
        if not self.lr_anneal_steps:
            return
        frac_done = (self.step + self.resume_step) / self.lr_anneal_steps
        lr = self.lr * (1 - frac_done)
        for param_group in self.opt.param_groups:
            param_group["lr"] = lr

    def log_step(self):
        """Log training step"""
        logger.logkv("step", self.step + self.resume_step)
        logger.logkv("samples", (self.step + self.resume_step + 1) * self.global_batch)

    def save(self):
        """Save model checkpoints"""
        def save_checkpoint(rate, params):
            state_dict = self.mp_trainer.master_params_to_state_dict(params)
            if (not dist.is_initialized() or dist.get_rank() == 0):
                logger.log(f"Saving model {rate}...")
                if not os.path.exists(self.resume_checkpoint):
                    logger.log(f"Creating directory: {self.resume_checkpoint}")
                    os.makedirs(self.resume_checkpoint, exist_ok=True)
                if not rate:
                    filename = f"model{(self.step+self.resume_step):06d}.pt"
                else:
                    filename = f"ema_{rate}_{(self.step+self.resume_step):06d}.pt"
                filepath = os.path.join(self.resume_checkpoint, filename)
                with open(filepath, "wb") as f:
                    th.save(state_dict, f)

        save_checkpoint(0, self.mp_trainer.master_params)
        for rate, params in zip(self.ema_rate, self.ema_params):
            save_checkpoint(rate, params)

        # Save optimizer state
        if (not dist.is_initialized() or dist.get_rank() == 0):
            opt_filename = f"opt{(self.step+self.resume_step):06d}.pt"
            opt_filepath = os.path.join(self.resume_checkpoint, opt_filename)
            with open(opt_filepath, "wb") as f:
                th.save(self.opt.state_dict(), f)

        if dist.is_initialized():
            dist.barrier()

def log_loss_dict(diffusion, ts, losses):
    """Log loss dictionary"""
    for key, values in losses.items():
        logger.logkv_mean(key, values.mean().item())
        # Log the quantiles (four quartiles, in particular).
        for sub_t, sub_loss in zip(ts.cpu().numpy(), values.detach().cpu().numpy()):
            quartile = int(4 * sub_t / diffusion.num_timesteps)
            logger.logkv_mean(f"{key}_q{quartile}", sub_loss)
